# 1 输出loss曲线
# 2 将记忆里的轨迹重新学习
from matplotlib import pyplot as plt 
import torch
import numpy as np
import logging
module_logger = logging.getLogger("mainModule.utils")

def plot(name, frame_idx, rewards):
    plt.figure(figsize=(20,5))
    plt.subplot(131)
    plt.title('frame %s. %s: %s' % (frame_idx, name, int(rewards[-1])))
    plt.plot(rewards)
    plt.savefig("%s.png" % (name))
    plt.close()

def v_wrap(np_array, dtype=np.float32):
    """
    作用：
    np_array -------->torch.tensor
    并且转换cuda
    """
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    if np_array.dtype != dtype:
        np_array = np_array.astype(dtype)
    return torch.from_numpy(np_array).to(device)

def set_init(layers):
    """
    初始化网络层
    :param layers:
    :return:
    """
    for layer in layers:
        nn.init.normal_(layer.weight, mean=0., std=0.1)
        nn.init.constant_(layer.bias, 0.)

def push_and_pull(opt, lnet, gnet, done, s_, bs, bat, bap, br, gamma):
    """
    optimize的更新，看不懂。
    """
    if done:
        v_s_ = 0.               # terminal
    else:
        v_s_ = lnet.forward(v_wrap(s_[None, :]))[-1].data.cpu().numpy()[0, 0]
    buffer_v_target = []
    for r in br[::-1]:    # reverse buffer r
        v_s_ = r + gamma * v_s_
        buffer_v_target.append(v_s_)
    buffer_v_target.reverse()

    loss = lnet.loss_func(
        v_wrap(np.vstack(bs)),
        v_wrap(np.array(bat), dtype=np.int64) if bat[0].dtype == np.int64 else v_wrap(np.vstack(bat)),
        v_wrap(np.vstack(bap)),
        v_wrap(np.array(buffer_v_target)[:, None]))

    # calculate local gradients and push local parameters to global
    opt.zero_grad()
    loss.backward()
    for lp, gp in zip(lnet.parameters(), gnet.parameters()):
        gp._grad = lp.grad
    opt.step()

    # pull global parameters
    lnet.load_state_dict(gnet.state_dict())

def record(global_ep, global_ep_r, ep_r, res_queue, name):
    with global_ep.get_lock():
        global_ep.value += 1
    with global_ep_r.get_lock():
        if global_ep_r.value == 0.:
            global_ep_r.value = ep_r
        else:
            global_ep_r.value = global_ep_r.value * 0.99 + ep_r * 0.01
    res_queue.put(global_ep_r.value)
    module_logger.info("Ep:" + str(global_ep.value) + "| Ep_r: %.0f" % global_ep_r.value)

def make_time_major(batch):
    """
    stack : 重叠维度
    transpose： 交换维度
    :param batch:
    :return:
    """
    states = []
    actions = []
    rewards = []
    dones = []
    next_states = []
    target_qs = []
    for t in batch:
        states.append(t.states)
        rewards.append(t.rewards)
        dones.append(t.dones)
        actions.append(t.actions)
        target_qs.append(t.target_qs)
        next_states.append(t.next_states)
    states = torch.stack(states).transpose(0, 1)
    actions = torch.stack(actions).transpose(0, 1)
    rewards = torch.stack(rewards).transpose(0, 1)
    dones = torch.stack(dones).transpose(0, 1)
    target_qs = torch.stack(target_qs).permute(0, 1)
    next_states = torch.stack(next_states).transpose(0, 1)
    return states, next_states, actions, rewards, target_qs, dones
